"""
API 速率限制器
确保 API 调用频率不超过限制
"""
import time
from collections import deque
from threading import Lock


class RateLimiter:
    """速率限制器"""

    def __init__(self, max_calls: int, time_window: int):
        """
        初始化速率限制器

        Args:
            max_calls: 时间窗口内最大调用次数
            time_window: 时间窗口（秒）
        """
        self.max_calls = max_calls
        self.time_window = time_window
        self.calls = deque()
        self.lock = Lock()

    def wait_if_needed(self):
        """如果需要，等待直到可以进行下一次调用"""
        with self.lock:
            now = time.time()

            # 移除时间窗口外的调用记录
            while self.calls and self.calls[0] < now - self.time_window:
                self.calls.popleft()

            # 如果达到限制，等待
            if len(self.calls) >= self.max_calls:
                sleep_time = self.calls[0] + self.time_window - now + 0.1
                if sleep_time > 0:
                    print(f"⏳ API 速率限制：等待 {sleep_time:.1f} 秒...")
                    time.sleep(sleep_time)
                    # 重新清理
                    now = time.time()
                    while self.calls and self.calls[0] < now - self.time_window:
                        self.calls.popleft()

            # 记录本次调用
            self.calls.append(time.time())

    def get_stats(self) -> dict:
        """获取统计信息"""
        with self.lock:
            now = time.time()
            # 清理过期记录
            while self.calls and self.calls[0] < now - self.time_window:
                self.calls.popleft()

            return {
                "current_calls": len(self.calls),
                "max_calls": self.max_calls,
                "time_window": self.time_window,
                "remaining": self.max_calls - len(self.calls)
            }


# 全局速率限制器：每分钟最多10次调用
api_rate_limiter = RateLimiter(max_calls=10, time_window=60)


def rate_limited_call(func):
    """
    装饰器：为函数添加速率限制

    Args:
        func: 要限制的函数

    Returns:
        包装后的函数
    """
    def wrapper(*args, **kwargs):
        api_rate_limiter.wait_if_needed()
        return func(*args, **kwargs)
    return wrapper


if __name__ == "__main__":
    # 测试速率限制器
    print("=" * 80)
    print("测试速率限制器")
    print("=" * 80)

    # 创建一个测试限制器：每5秒最多3次
    test_limiter = RateLimiter(max_calls=3, time_window=5)
    
    print("\n测试：快速调用5次（限制：5秒内最多3次）")
    for i in range(5):
        test_limiter.wait_if_needed()
        print(f"调用 {i+1}: {time.strftime('%H:%M:%S')}")
        stats = test_limiter.get_stats()
        print(f"  当前: {stats['current_calls']}/{stats['max_calls']}, 剩余: {stats['remaining']}")
